import cv2
from pathlib import Path
from utiils.utills_functions import readConf
from utiils.yolov5_detector import Person_YOLO5_Detector
VID_FORMATS = 'asf', 'avi', 'gif', 'm4v', 'mkv', 'mov', 'mp4', 'mpeg', 'mpg', 'ts', 'wmv'
import os

def img_predector(frame, Person_detector):
    pred, names = Person_detector(frame)
    for i, box in enumerate(pred):
        cv2.rectangle(frame, (int(box[0]), int(box[1])),
                      (int(box[2]), int(box[3])),
                      (0, 255, 0), 2)

        cv2.putText(frame, str(i), (int(box[0]), int(box[1])),
                    cv2.FONT_HERSHEY_PLAIN, 2,
                    (0, 0, 255), 2)
    r_humanCount = 'num: %d' % (len(pred))
    cv2.putText(frame, r_humanCount, (int(25), int(25)),
                cv2.FONT_HERSHEY_PLAIN, 2,
                (0, 0, 255), 2)
    return frame, r_humanCount

def Person_detector_run(inputs, output, conf):
    global frame, box
    frameNumber = 5
    r_humanCount = 0
    Person_detector = Person_YOLO5_Detector(conf)
    if Path(inputs).suffix[1:] in (VID_FORMATS):
        cap = cv2.VideoCapture()
        cap.open(inputs)
        fps = cap.get(cv2.CAP_PROP_FPS)
        width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)  # float
        height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
        fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
        out = cv2.VideoWriter(output, fourcc, fps, (int(width), int(height)))
        count = 0
        while cap.grab():
            count = count + 1
            # 帧数从1开始统计，将25帧转为24帧，
            if count % frameNumber == 0:
                print("正在处理第{}帧".format(count))
                ret, frame = cap.retrieve()
                image = frame[:, :, [2, 1, 0]]
                if not ret:
                    continue
                # 检测框 pre :torch.Size([2, 6])  tensor([[540.00000, 245.00000, 753.00000, 800.00000,   0.93690,   2.00000]  [646.00000, 397.00000, 719.00000, 775.00000,   0.89301,   0.00000]])
                frame, human_count = img_predector(frame, Person_detector)
                out.write(frame)

    elif os.path.isdir(inputs):
        for filename in os.listdir(inputs):
            if filename.endswith(".jpg") or filename.endswith(".png"):
                img_path = os.path.join(inputs, filename)
                frame = cv2.imread(img_path)
                if frame is None:
                    print(f"无法读取图片: {img_path}")
                else:
                    frame, human_count = img_predector(frame, Person_detector)
                    print(human_count)
                    cv2.imwrite('./' + output + '/' + filename[:-3] + "jpg", frame)

    else:
        frame = cv2.imread(inputs)
        frame, human_count = img_predector(frame, Person_detector)
        cv2.imwrite(output[:-3] + "jpg", frame)
    r_humanCount = human_count

    return r_humanCount

if __name__ == '__main__':
    inputs = './test_datasets/input_datasets'  # 可输入、视频、图片、图片文件夹路径
    output = "test_datasets/output_datasets"  # 指定存放的文件夹路径
    conf = readConf("config_set/model.cfg")
    Person_detector_run(inputs, output, conf)
